package com.pai4j.search.controller;

import com.pai4j.common.bean.PAIResponseBean;
import com.pai4j.common.bean.PAIResponseBeanUtil;
import com.pai4j.domain.vo.request.search.MilvusSearchRequestVO;
import com.pai4j.log.annotation.OperationLog;
import com.pai4j.remote.search.VectorSearchEngineClient;
import com.pai4j.search.service.voctor.AbstractMilvusService;
import com.pai4j.search.service.voctor.MilvusServiceFactory;
import org.springframework.util.Assert;
import org.springframework.web.bind.annotation.RestController;
/**
 *  向量搜索引擎服务
 *
 * @author: CYM-pai
 * @date: 2025/05/29 11:22
 **/

@RestController
public class VectorSearchEngineController implements VectorSearchEngineClient {

    /**
     *  根据对应的场景创建向量集合
     *
     * @author: CYM-pai
     * @date: 2025/05/29 11:23
     **/
    @Override
    @OperationLog(
        module = "向量搜索",
        operation = "创建向量集合",
        description = "管理员创建向量搜索集合",
        recordParams = true
    )
    public PAIResponseBean<Boolean> createCollection(String biz) {
        MilvusServiceFactory.getService(biz).createCollection();
        return PAIResponseBeanUtil.success(true);
    }

    
    /**
     *
     *  根据对应的场景删除向量集合
     * @author: CYM-pai
     * @date: 2025/05/29 11:23
     **/
    @Override
    public PAIResponseBean<Boolean> dropCollection(String biz) {
        MilvusServiceFactory.getService(biz).dropCollection();
        return PAIResponseBeanUtil.success(true);
    }

    
    /**
     *
     *  查询向量集合
     * @author: CYM-pai
     * @date: 2025/05/29 11:26
     **/
    @Override
    @OperationLog(
        module = "向量搜索",
        operation = "向量搜索",
        description = "用户进行向量相似度搜索",
        recordParams = true
    )
    public PAIResponseBean<Object> search(String biz, MilvusSearchRequestVO milvusSearchReq) {
        String query = milvusSearchReq.getQuery();
        Assert.notNull(query, "query must not be null!");
        Integer topK = milvusSearchReq.getTopK();
        //  topK 默认返回一条数据
        topK = topK == null ? 1 : topK;
        float score = milvusSearchReq.getScore();
        AbstractMilvusService milvusService = MilvusServiceFactory.getService(biz);
        Object result = milvusService.search(query, topK, score);
        return PAIResponseBeanUtil.success(result);
    }
}
